#!/usr/bin/env python
import copy
import re

from Ganga.GPIDev.Lib.File import File
from GangaAtlas.Lib.AMIGA.AMIDataset import AMIDataset
from AutoD3PDMakerTaskGenerator import AutoD3PDMakerTaskGenerator

class AMAEventMakerTaskGenerator(AutoD3PDMakerTaskGenerator):

    def __init__(self):
        AutoD3PDMakerTaskGenerator.__init__(self)

    def getDatasetsFromGRL(self, grlPath, streams, runTags):

        ds = AMIDataset()
        ds.goodRunListXML = File(grlPath)

        dlist_tmp = ds.search()

        datasets = {}
        for ds in dlist_tmp:
            info  = ds.split('.')
            prj   = info[0]
            rid   = info[1]
            strm  = info[2]
            rtag  = re.sub('\/$','', info[-1] )
            dtype = info[-2]

            if rtag not in runTags:
                continue

            if strm not in streams:
                continue

            if re.match('^data10.*',prj) and re.match('^f[0-9]+.*', rtag):
                datasets['%s_%s_%s' % (prj,strm,rid)] = [ re.sub('\/$','',ds) ]
            else:
                datasets['%s_%s_%s' % (prj,strm,rid)] = [ ds ]

        return datasets

    def getTaskTemplate(self, args_dict={}):

        taskTemplate = {'ama_userarea' : 'amaathena_userarea-00-01-81.tar.gz',
                        'amaconfig'    : 'AMAEvents_mc09_7TeV_Common_SignalMcAtNlo.py',
                        'input'        : 'mc09_7TeV.105200.T1_McAtNlo_Jimmy.merge.AOD.e510_s624_s633_r1085_r1113/',
                        'metadata'     : {'version'       : 1,
                                          'sample'        : 'mc09_7TeV',
                                          'signal'        : 'TTbar_MCatNLO_NoAllHadronic',
                                          'run'           : '105200',
                                          'algorithm'     : 'T1_McAtNlo_Jimmy',
                                          'tag'           : 'e510_s624_s633_r1085_r1113',
                                          'atlas_release' : '15.6.9',
                                          'atlas_project' : '',
                                          'nickname'      : '',
                                          'ama_flags'     : ['MC09','IGNORETAGDIFF','TOPINPUTS','TRIG']}}

        ## override template values if the key:value pair is given in the args_dict
        for k, v in args_dict.items():
            if k in taskTemplate.keys():

                if type(v) == type(dict()):
                    taskTemplate[k].update(v)
                else:
                    taskTemplate[k] = v

            else:
                taskTemplate[k] = v
                
        return taskTemplate

    def generateTasks(self, template='', args_dict={}):
        """
        generates AMAEvent generation tasks.

        The input to generate AMAEvent task is basically a list of dataset.
        The dataset list can be given using one of the following two options:

        option 1: given a list of dataset manually

                  e.g. args_dict['datasets'] = {'data10_physics_MinBias_00152409': ['data10_7TeV.00152409.physics_MinBias.merge.AOD.f238_m427'],
                                                'data10_physics_MinBias_00152508': ['data10_7TeV.00152508.physics_MinBias.merge.AOD.f241_m433']}

        option 2: given a GRL.xml (together with some optional attributes) for dataset selection

                  e.g. args_dict['grlist'] = {'xml'     : '/localstore/hclee/autod3pdmaker/grl/GRL_20100531_lepton_streams.xml',
                                              'streams' : ['physics_MuonswBeam','physics_L1Calo'],
                                              'runtags' : ['r1297_p161','f259_m486','f260_m496']}

                  the 'xml' attribute is mandatory; while 'streams' and 'runtags' attributes are optional.

        if both option 1 and option 2 are specified, all datasets will be taken into account in generating tasks.
        """

        if not template:
            template = self.getTaskTemplate()

        amaconfigMap = { 'TTbar_Powheg'  : 'AMAEvents_mc09_7TeV_Common_SignalPowHeg.py',
                         'TTbar_MCatNLO' : 'AMAEvents_mc09_7TeV_Common_SignalMcAtNlo.py',
                         'background'    : 'AMAEvents_mc09_7TeV_Common_Bkg.py',
                         'qcd_background': 'AMAEvents_mc09_7TeV_Common_QCD_Bkg.py',
                         'data10'        : 'AMAEvents_data10_7TeV_Common.py',
                         'group'         : 'AMAEvents_data10_7TeV_Common.py',
                         'MC_MinBias'    : 'AMAEvents_data10_7TeV_Common.py' }


        dsDict = {}

        try:
            dsDict.update( args_dict['datasets'] )
        except KeyError:
            pass

        try:
            grlInfo = args_dict['grlist']
            grlPath = grlInfo['xml']

            streams = []
            if grlInfo.has_key('streams'):
                streams = grlInfo['streams']

            runTags = []
            if grlInfo.has_key('runtags'):
                runTags = grlInfo['runtags']

            dsDict.update( self.getDatasetsFromGRL(grlPath, streams, runTags) )
        except KeyError:
            pass

        if not dsDict:
            self.logger.warning('no dataset list for generating AMAEventMaker tasks')

        tlist = []

        for prj,dsList in dsDict.items():

            amaconfig = amaconfigMap['background']
            amaflags  = ['MC09','IGNORETAGDIFF','RUNJETSMC','TOPINPUTS','TRIG']
            if re.match('TTbar_PowHeg', prj, re.I):
                amaconfig = amaconfigMap['TTbar_Powheg']
            elif re.match('TTbar_MCatNLO', prj, re.I):
                amaconfig = amaconfigMap['TTbar_MCatNLO']
            elif re.match('MC_MinBias', prj, re.I):
                amaconfig = amaconfigMap['MC_MinBias']
            elif re.match(r'.*QCD.*', prj, re.I):
                amaconfig = amaconfigMap['qcd_background']
            elif re.match('data10', prj, re.I):
                amaconfig = amaconfigMap['data10']
                amaflags  = ['DATA','RUNJETSDATA','TOPINPUTS','TRIG']
            elif re.match('group', prj, re.I):
                amaconfig = amaconfigMap['group']
                amaflags  = ['DATA','IGNORETAGDIFF','TOPINPUTS','TRIG']

            for ds in dsList:

                if ds:
                    t = copy.deepcopy(template)
                    t['input']              = ds
                    t['amaconfig']          = amaconfig
                    t['metadata']['signal'] = prj

                    ds_attrs = re.sub('\/$','',ds).split('.')
                    t['metadata']['sample']    = ds_attrs[0]
                    t['metadata']['run']       = ds_attrs[1]
                    t['metadata']['algorithm'] = ds_attrs[2]
                    t['metadata']['tag']       = ds_attrs[-1]
                    t['metadata']['ama_flags'] = amaflags

                    ## something adhoc only for 2010 May reprocessing
                    if re.match('^data10.*\.f[0-9]+.*', ds):
                        t['metadata']['nickname'] = 't0proc04'

                    if re.match('^data10.*\.r[0-9]+.*', ds):
                        t['metadata']['nickname'] = 'repro04'

                    tlist.append(t)

        return tlist
